'''下载一个m3u8的hls流'''

import sys
import re
import time

import requests
from requests.adapters import HTTPAdapter

class m3u8Downloader:
  
  ## 传入m3u8地址,需要下载的文件地址,m3u8的初始下载文件(可以在中途开始),
  ## 写入缓冲大小(默认20MB)
  def __init__(self, m3u8, filePath, startIndex = 0,
    bufferSize = 1024*1024*20):
    self.bufferSize = bufferSize
    self.startIndex = startIndex
    self.filePath = filePath
    self.links = [] # 下载的链接表
    self.session = requests.Session()
    self.session.mount('http://', HTTPAdapter(max_retries=3))
    self.session.mount('https://', HTTPAdapter(max_retries=3))
    self._handleM3u8(m3u8) # 处理m3u8,获取下载链接
    print('Ready to Download {0}\n, parts ='.format(self.filePath), 
      len(self.links))
    return
    
  ## 处理m3u8,里面的条目如果是m3u8,则递归处理,如果是ts等视频片段,
  ## 则加入下载列表
  def _handleM3u8(self, m3u8):
    urlTop = re.sub(r'[^/]*$', '', m3u8)
    resp = self.session.get(m3u8)
    texts = filter(lambda x:(not (x.startswith('#'))) and 
      (len(x.strip()) > 0), resp.text.split('\n'))
    for url in texts:
      if url.endswith('.m3u8'):
        self._handleM3u8(urlTop+url)
      else:
        self.links.append(urlTop+url)
    return
    
  ## 主函数,下载用.如stepConfirm是True,每写入一次缓存要确认一次.
  ## lastIndex指示下一次的下载位置,每写入一次文件更新一次,若中途出问题,记住
  ## 这个显示的数字,下次指定从这里开始下载就行了.
  def main(self, stepConfirm = False):
    tempContent = b'' # 下载的内容的总和,每到缓存前清空一次
    lastIndex = self.startIndex
    print('\nStart Download from index '+str(lastIndex))  
    timeStart = time.time() # 当前时间,计算下载速度,每写入文件一次计算一次
    for i in range(self.startIndex, len(self.links)):
      try:
        link = self.links[i]
        resp = self.session.get(link)
        if resp.status_code != 200:
          print('Error!'+link+' status_code='+str(resp.status_code))
          print('lastIndex='+str(lastIndex))
          return
        tempContent += resp.content
        # 显示未保存的下载量
        mb = len(tempContent)/1048576.0
        sys.stdout.write('\rGot {0:.2f}MB'.format(mb))
        sys.stdout.flush()
        ## 缓存大小够了就写入文件
        if len(tempContent) > self.bufferSize:
          with open(self.filePath, 'ab+') as f:
            f.write(tempContent)
          now = time.time()
          speedKBS = len(tempContent)/(1024.0)/(now-timeStart+0.0001)
          print('\nWrite buffer to file,download speed={0:.2f}KB/s'.format(speedKBS))
          lastIndex = i+1
          tempContent = b''
          if stepConfirm:
            print('\nlastIndex = '+str(lastIndex))
            confirmMsg = 'Part of file written,press Enter to continue or '
            confirmMsg += '11Enter to exit'
            x = input(confirmMsg)
            if x == '11':
              return
          timeStart = time.time()
      except Exception as e:
        print('Error!'+link+' exception:',e)
        print('lastIndex='+str(lastIndex))
        return
    with open(self.filePath, 'ab+') as f:
      f.write(tempContent)
    now = time.time()
    speedKBS = len(tempContent)/(1024.0)/(now-timeStart+0.001)
    print('\nWrite content to file,speed={0:.2f}KB/s'.format(speedKBS))
    print('Download OK!')
    return
  
      
if __name__ == '__main__':
  m3u8 = r'https://youku.com-youku.net/20180526/6375_bd55a423/index.m3u8'
  filePath = r'F:\Download\视频下载\测试\大唐双龙传第16集.mp4'
  startIndex = 0
  bufferSize = 1024*1024*50
  
  md = m3u8Downloader(m3u8, filePath, startIndex = startIndex, 
    bufferSize = bufferSize)
  input('Class m3u8Downloader init finished!Press Enter to continue')
  md.main(stepConfirm = False) # stepConfirm设定为False,则一直下载不需要确认